# encoding: utf-8
"""
通过query data 获取数据
"""
import json
from typing import Optional, Literal, List, Dict, Any

from pydantic import BaseModel, Field


# 定义的action流程
class ActionFlowDataBase(BaseModel):
    name: str = Field(description="流程的名字")
    flowSuccess: bool = Field(description="是否成功")


class ErrorActionFlowData(ActionFlowDataBase):
    flowSuccess: bool = False
    errorMsg: str = Field(description="错误信息")


class QueryFlowData(ActionFlowDataBase):
    flowSuccess: bool = True
    originQuery: str = Field(description="原始提问")
    rewriteQuery: str = Field(description="改写提问")
    dataExists: bool = Field(description="是否有数据")
    dataLength: int = Field(description="获取的数据长度")


class SupplyDataBase(BaseModel):
    type: str = Field(default='Any', description='数据类型')
    description: str = Field(default='', description='数据描述')
    data: Any = Field(default=None, description='数据')


class TableData(BaseModel):
    query: str = Field('', description='智能问数时的查询语句')
    condition: str = Field(default='', description='数据表的条件')
    fields: Any = Field(default='', description='数据表字段的简略说明')
    columns: List = Field(default=[], description='数据表字段的详细说明')
    rows: List[Dict[str, Any]] = Field(default=[], description='行数据')


class AiQueryData(SupplyDataBase):
    type: str = 'TableData'
    description: str = Field(default=None, description='数据的描述')
    data: TableData = Field(description='智能问数查询的表格数据')


class EmptyData(SupplyDataBase):
    type: str = 'EmptyData'
    description: str = Field(description='数据的描述')
    data: None = None


class UploadData(SupplyDataBase):
    """上传文件的数据"""
    type: str = 'Any'
    description: str = Field(default='', description='用户上传数据的数据描述')
    data: Any = Field(default=None, description='用户上传数据')


class UserInfoData(SupplyDataBase):
    """用户信息数据"""
    type: str = 'Any'
    description: str = Field(default='', description='用户信息描述')
    data: Any = Field(default=None, description='用户数据')


class WebSearchData(SupplyDataBase):
    """网络检索数据"""
    type: str = 'Any'
    description: str = Field(default='', description='网络检索数据描述')
    data: Any = Field(default=None, description='网络检索数据')


class Response(BaseModel):
    msg: Literal['success', 'error'] = Field(default='success', description='是否成功返回结果')
    data: List[UserInfoData | AiQueryData | WebSearchData | UploadData] = Field(default=[], description='数据列表')
    error: str | None = Field(default=None, description='错误说明')


def _parse_query_data(data: str | Dict) -> SupplyDataBase:
    if isinstance(data, str):
        data = json.loads(data)

    dtype = data['type']
    print('data search dtype:', dtype, data['data'] is None)
    if data['data'] is None:
        return EmptyData(
            description=data['description'],
        )

    if dtype == 'TableData':
        return AiQueryData(**data)

    else:
        raise ValueError(f'未能识别的type: {dtype}')


def parse_query_data(data: str | List | Dict, filter_empty: bool = True) -> List[SupplyDataBase]:
    if isinstance(data, str):
        data = json.loads(data)

    if isinstance(data, Dict):
        data = List[data]

    result = []
    for _ in data:
        new_one = _parse_query_data(_)
        if isinstance(new_one, EmptyData):
            if filter_empty:
                continue
            else:
                result.append(new_one)
        else:
            result.append(new_one)
    return result


def parse_flow_data(origin_query: str, supply_data: List[SupplyDataBase]) -> List[ActionFlowDataBase]:
    """
    封装数据解析
    :param origin_query: 原始提问
    :param supply_data:
    :return:
    """
    all_flow_data = []
    for data in supply_data:
        if isinstance(data, EmptyData):
            flow_data = QueryFlowData(
                name='wenshu',
                originQuery=origin_query,
                rewriteQuery=data.description,
                dataExists=False,
                dataLength=0
            )
        elif isinstance(data, AiQueryData):
            flow_data = QueryFlowData(
                name='wenshu',
                originQuery=origin_query,
                rewriteQuery=data.data.query,
                dataExists=True,
                dataLength=len(data.data.rows)
            )
        else:
            continue
        all_flow_data.append(flow_data)
    return all_flow_data